from datasets import load_dataset, load_from_disk
from torch.utils.data import Dataset
# pip install datasets -i https://pypi.tuna.tsinghua.edu.cn/simple

#在线加载数据集或者本地加载数据集
cach_dir = "dataset"
# load_dataset(path="lansinuote/ChnSentiCorp", cach_dir=cach_dir)
dataset = load_from_disk(r"D:\pythonWork\python_demo\dataset\ChnSentiCorp")
#打印出数据集的具体内容
test_data = dataset["test"]
print(dataset["test"])
for i in test_data:
    print(i)
#继承pytoch Dataset重写三个方法
class MyDataset(Dataset):
    # 初始化数据集
    def __init__(self, split):
        self.dataset = load_from_disk(r"D:\pythonWork\PythonProject3\dataset\ChnSentiCorp")
        if split == "train":
            # 训练集
            self.dataset = self.dataset["train"]
        elif split == "test":
            # 测试集
            self.dataset = self.dataset["test"]
        elif split == "validation":
            # 验证集
            self.dataset = self.dataset["validation"]
        else:
            print("split must be train, test or validation")
    #对数据做定制化处理
    def __getitem__(self, index):
        text =  self.dataset[index]['text']
        label =  self.dataset[index]['label']
        # 样例数据 {'text': '为了高兴我买了第二个小白，这个小白安静，不爱叫，外皮漂亮，得常擦擦，长的也紧凑，有很好的设计感，虽然个子小，什么都安得上，装了不少东西，跑的一点不慢', 'label': 1}
        return text,label
    #返回数据集长度
    def __len__(self):
        return len(self.dataset)
if __name__ == '__main__':
    # 创建数据集对象
    data = MyDataset("validation")
    for i in data:
        print(i)